package ai.mspbots.sql.service;

import ai.mspbots.sql.datasource.DataSourceContextHolder;
import ai.mspbots.sql.datasource.DynamicDataSource;
import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.List;
import java.util.concurrent.CountDownLatch;

@Service
@Slf4j
public class SQLParserService {

    @Autowired
    private DataSource dataSource; // 这里注入的是 DynamicRoutingDataSource

    @Autowired
    private ThreadPoolTaskExecutor taskExecutor;

    /**
     * 获取当前数据源的 DataSource
     */
    public DataSource getCurrentDataSource() {
        // 获取当前数据源的 key
        String dataSourceKey = DynamicDataSourceContextHolder.peek();
        if (dataSourceKey == null) {
            throw new IllegalStateException("当前没有指定数据源");
        }

        // 强制转换为 DynamicRoutingDataSource
        if (dataSource instanceof DynamicRoutingDataSource) {
            DynamicRoutingDataSource dynamicDataSource = (DynamicRoutingDataSource) dataSource;
            return dynamicDataSource.getDataSource(dataSourceKey);
        } else {
            throw new IllegalStateException("数据源不是 DynamicRoutingDataSource 类型");
        }
    }

    public String assemble(String tableName, String selectQuery, String countQuery, String date) {
        DataSource dataSource = getCurrentDataSource();
        String base = "C:\\sql\\" + tableName + "-{}.txt";
        if (!StringUtils.hasText(date)) {
            date = DateUtil.date().toString(DatePattern.PURE_DATETIME_PATTERN);
        }
        String path = StrUtil.format(base, date);

        long start = System.currentTimeMillis();
        log.info("开始执行, tableName: {}, path: {}", tableName, path);
        File file = new File(path);

        try {
            if (file.exists()) {
                boolean delete = file.delete();
                log.debug("delete: {}", delete);
            }
            boolean newFile = file.createNewFile();
            log.debug("newFile: {}", newFile);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        try (Connection connection = dataSource.getConnection()) {
            // 查询总条数
            PreparedStatement countStmt = connection.prepareStatement(countQuery);
            ResultSet countResultSet = countStmt.executeQuery();
            int totalRecords = 0;
            if (countResultSet.next()) {
                totalRecords = countResultSet.getInt(1);
            }
            countResultSet.close();
            countStmt.close();

            int pageSize = 200;  // 每页的记录数
            int totalPages = (int) Math.ceil((double) totalRecords / pageSize);
            CountDownLatch latch = new CountDownLatch(totalPages);

            for (int pageNumber = 0; pageNumber < totalPages; pageNumber++) {
                int finalPageNumber = pageNumber; // 用于闭包
                String datasource = DynamicDataSourceContextHolder.peek();
                taskExecutor.execute(() -> {
                    DynamicDataSourceContextHolder.push(datasource);
                    try {
                        writePageToFile(tableName, selectQuery, finalPageNumber, pageSize, path);
                    } finally {
                        latch.countDown();
                        DynamicDataSourceContextHolder.clear();
                    }
                });
            }

            // 等待所有线程完成
            latch.await();
            log.info("结束执行, 写入文件: {}, 共: {} 条数据, 耗时: {}s",path, totalRecords, (System.currentTimeMillis() - start) / 1000);
            return "Completed writing to " + path;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    private void writePageToFile(String tableName, String selectQuery, int pageNumber, int pageSize, String path) {
        String paginatedQuery = selectQuery + " LIMIT " + pageSize + " OFFSET " + (pageNumber * pageSize);
        DataSource dataSource = getCurrentDataSource();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement selectStmt = connection.prepareStatement(paginatedQuery);
             ResultSet resultSet = selectStmt.executeQuery()) {

            StringBuilder result = new StringBuilder();
            while (resultSet.next()) {
                StringBuilder insertSql = new StringBuilder();
                StringBuilder deleteSql = new StringBuilder();

                long id = resultSet.getLong("id");
                deleteSql.append("DELETE FROM ").append(tableName).append(" WHERE id = ").append(id).append(";");
                insertSql.append("INSERT INTO ").append(tableName).append("(");

                int columnCount = resultSet.getMetaData().getColumnCount();
                for (int i = 1; i <= columnCount; i++) {
                    insertSql.append(resultSet.getMetaData().getColumnName(i));
                    if (i < columnCount) {
                        insertSql.append(", ");
                    }
                }
                insertSql.append(") VALUES (");

                for (int i = 1; i <= columnCount; i++) {
                    Object value = resultSet.getObject(i);
                    String columnType = resultSet.getMetaData().getColumnTypeName(i);

                    if (value == null) {
                        insertSql.append("NULL");
                    } else {
                        String formattedValue = getDbType(columnType, value);
                        insertSql.append(formattedValue);
                    }

                    if (i < columnCount) {
                        insertSql.append(", ");
                    }
                }
                insertSql.append(");");

                result.append(deleteSql).append("\n").append(insertSql).append("\n");
            }

            // 将结果写入文件
            FileUtil.appendLines(List.of(result.toString()), path, "UTF-8");
        } catch (Exception e) {
            log.error("Error writing page {} to file: {}", pageNumber, e.getMessage());
        }
    }


    private static String getDbType(String dbType, Object value) {
        switch (dbType) {
            case "bigserial":
            case "int8":
            case "integer":
                return value.toString(); // 数字类型不需要引号
            case "timestamp":
            case "varchar":
                return "'" + value.toString().replace("'", "''") + "'"; // 字符串类型需要转义
            case "jsonb":
                return "'" + value.toString().replace("'", "''") + "'::jsonb"; // JSONB 类型转义
            case "boolean":
                return (Boolean) value ? "true" : "false"; // 布尔类型处理
            default:
                return "'" + value.toString().replace("'", "''") + "'"; // 默认转义
        }
    }
}